# coding: utf-8
"""
Magnetostatic Fields
=====================

An example of using PlasmaPy's `Magnetostatic` class in `physics` subpackage.
"""

import plasmapy as pp
from plasmapy.physics import magnetostatics
from plasmapy.classes.sources import Plasma3D
import numpy as np
import astropy.units as u
import matplotlib.pyplot as plt
import itertools

############################################################
# Some common magnetostatic fields can be generated and added to a plasma object.
# A dipole

dipole = magnetostatics.MagneticDipole(np.array([0, 0, 1])*u.A*u.m*u.m, np.array([0, 0, 0])*u.m)
print(dipole)

############################################################
# initialize a a plasma, where the magnetic field will be calculated on

plasma = Plasma3D(domain_x=np.linspace(-2, 2, 30) * u.m,
                domain_y=np.linspace(0, 0, 1) * u.m,
                domain_z=np.linspace(-2, 2, 20) * u.m)

############################################################
# add the dipole field to it
plasma.add_magnetostatic(dipole)

X, Z = plasma.grid[0, :, 0, :], plasma.grid[2, :, 0, :]
U = plasma.magnetic_field[0, :, 0, :].value.T  # because grid uses 'ij' indexing
W = plasma.magnetic_field[2, :, 0, :].value.T  # because grid uses 'ij' indexing


############################################################
plt.figure()
plt.axis('square')
plt.xlim(-2, 2)
plt.ylim(-2, 2)
plt.title('Dipole field in x-z plane, generated by a dipole pointing in the z direction')
plt.streamplot(plasma.x.value, plasma.z.value, U, W)

############################################################
# A circular current-carring wire

cw = magnetostatics.CircularWire(np.array([0, 0, 1]), np.array([0, 0, 0])*u.m, 1*u.m, 1*u.A)
print(cw)

############################################################
# initialize a a plasma, where the magnetic field will be calculated on
plasma = Plasma3D(domain_x=np.linspace(-2, 2, 30) * u.m,
                domain_y=np.linspace(0, 0, 1) * u.m,
                domain_z=np.linspace(-2, 2, 20) * u.m)

############################################################
# add the circular coil field to it
plasma.add_magnetostatic(cw)

X, Z = plasma.grid[0, :, 0, :], plasma.grid[2, :, 0, :]
U = plasma.magnetic_field[0, :, 0, :].value.T  # because grid uses 'ij' indexing
W = plasma.magnetic_field[2, :, 0, :].value.T  # because grid uses 'ij' indexing

plt.figure()
plt.axis('square')
plt.xlim(-2, 2)
plt.ylim(-2, 2)
plt.title('Circular coil field in x-z plane, generated by a circular coil in the x-y plane')
plt.streamplot(plasma.x.value, plasma.z.value, U, W)

############################################################
# a circular wire can be described as parametric equation and converted to GeneralWire

gw_cw = cw.to_GeneralWire()

# the calculated magnetic field is close
print(gw_cw.magnetic_field([0, 0, 0]) - cw.magnetic_field([0, 0, 0]))


############################################################
# A infinite straight wire


iw = magnetostatics.InfiniteStraightWire(np.array([0, 1, 0]), np.array([0, 0, 0])*u.m, 1*u.A)
print(iw)

############################################################
# initialize a a plasma, where the magnetic field will be calculated on
plasma = Plasma3D(domain_x=np.linspace(-2, 2, 30) * u.m,
                domain_y=np.linspace(0, 0, 1) * u.m,
                domain_z=np.linspace(-2, 2, 20) * u.m)

# add the infinite straight wire field to it
plasma.add_magnetostatic(iw)

X, Z = plasma.grid[0, :, 0, :], plasma.grid[2, :, 0, :]
U = plasma.magnetic_field[0, :, 0, :].value.T  # because grid uses 'ij' indexing
W = plasma.magnetic_field[2, :, 0, :].value.T  # because grid uses 'ij' indexing

plt.figure()
plt.title('Dipole field in x-z plane, generated by a infinite straight wire '
          'pointing in the y direction')
plt.axis('square')
plt.xlim(-2, 2)
plt.ylim(-2, 2)
plt.streamplot(plasma.x.value, plasma.z.value, U, W)
